#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
供水调度优化求解器模块
基于pyoptinterface的多求解器支持
"""

from loguru import logger
from typing import Dict, Any, Optional
import pyoptinterface as poi
from pyoptinterface import mosek, gurobi, highs, copt
from optmodel.log_config import timer


def get_solver(params: Dict[str, Any]) -> object:
    """
    Retrieve the solver object based on parameters.
    
    Parameters
    ----------
    params : dict
        Configuration dictionary with solver details.
        
    Returns
    -------
    object
        Solver object based on the configuration. 
    """
    solver_map = {
        'mosek': mosek,
        'gurobi': gurobi, 
        'highs': highs,
        'copt': copt
    }
    
    solver_name = params.get('solver', {}).get('name', 'highs')
    
    if solver_name in solver_map:
        poi_solver = solver_map[solver_name]
    else:
        raise ValueError(f"不支持的求解器: {solver_name}")
        
    # 尝试自动加载求解器库
    if not poi_solver.autoload_library():
        logger.warning(f"{solver_name} 库自动加载失败，尝试手动加载")
        
        if 'path' not in params.get('solver', {}):
            raise ValueError(f"配置中未找到 {solver_name} 的求解器路径")
            
        solver_path = params['solver']['path']
        if not poi_solver.load_library(solver_path):
            raise RuntimeError(f"无法从路径 {solver_path} 加载 {solver_name} 库")
            
    logger.info(f"成功加载求解器: {solver_name}")
    return poi_solver

def set_solver_parameters(model : object) -> None:
    """Set the solver-specific parameters for the model.

    Parameters
    ----------
    model : object
        Model to configurable.
    """
    # 获取求解器名称
    solver_name = model.param.get('solver', {}).get('name', 'highs')
    
    # 设置求解器日志输出
    try:
        if solver_name.lower() == 'highs':
            # HiGHS 求解器日志配置
            model.set_raw_parameter("log_to_console", True)
            model.set_raw_parameter("output_flag", True)
            model.set_raw_parameter("log_dev_level", 1)  # 详细日志级别
            
        elif solver_name.lower() == 'gurobi':
            # Gurobi 求解器日志配置
            model.set_raw_parameter("OutputFlag", 1)
            model.set_raw_parameter("LogToConsole", 1)
            model.set_raw_parameter("DisplayInterval", 1)
            
        elif solver_name.lower() == 'mosek':
            # MOSEK 求解器日志配置
            model.set_raw_parameter("MSK_IPAR_LOG", 1)
            model.set_raw_parameter("MSK_IPAR_LOG_LEVEL", 4)
            
        elif solver_name.lower() == 'copt':
            # COPT 求解器日志配置
            model.set_raw_parameter("Logging", 1)
            model.set_raw_parameter("LogToConsole", 1)
            
        logger.info(f"已配置 {solver_name} 求解器日志输出")
        
    except Exception as e:
        logger.warning(f"配置求解器日志参数时出错: {e}")
    
    # 设置其他用户自定义参数
    for key, value in model.param.get('solver', {}).items():
        if key not in ('name', 'path'):
            try:
                model.set_raw_parameter(key, value)
                logger.debug(f"设置求解器参数: {key} = {value}")
            except Exception as e:
                logger.warning(f"设置参数 {key} 失败: {e}")

@timer
def solve_model(model: object) -> bool:
    """
    Solve the model.
    
    Parameters
    ----------
    model : object
        Model to solve.
        
    Returns
    -------
    bool
        True if the model is solved optimally, False otherwise.
    """
    model.optimize()
    status = model.get_model_attribute(poi.ModelAttribute.TerminationStatus)
    if status != poi.TerminationStatusCode.OPTIMAL:
        return False
    return True
